#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import re
import argparse
import logging
from PyPDF2 import PdfReader, PdfWriter

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def extract_heading_candidates(pdf_path):
    """从PDF文件中提取可能的标题内容"""
    try:
        logger.info(f"分析PDF文件: {pdf_path}")
        reader = PdfReader(pdf_path)
        headings = []
        
        # 遍历每一页查找可能的标题
        for page_num, page in enumerate(reader.pages, 1):
            text = page.extract_text()
            if not text:
                continue
                
            lines = text.split('\n')
            for line in lines:
                line = re.sub(r'\s+', '', line)
                # line = line.strip()
                # 查找可能是标题的文本行：短小、以数字或特定字符开头
                if 2 <= len(line) <= 100:
                    # 标题模式1: 数字开头（如 "1.xxx"）
                    if re.match(r'^\d+(\.\d+)*\.', line):
                        print(line)
                        headings.append((line, page_num))
                    # 标题模式2: "第X章"、"第X节"等
                    elif re.match(r'^[第章节篇部分]\s*[一二三四五六七八九十\d]+', line):
                        headings.append((line, page_num))
                    # 标题模式3: 全大写或首字母大写的短句
                    # elif line.isupper() or (len(line.split()) <= 7 and line[0].isupper()):
                    #     headings.append((line, page_num))
        
        return headings
    except Exception as e:
        logger.error(f"分析文件 {pdf_path} 时出错: {str(e)}")
        logger.exception(e)
        return []

def add_bookmarks_to_pdf(pdf_path, output_path=None, auto_detect=True, force_detect=False):
    """向PDF添加书签"""
    if not output_path:
        filename, ext = os.path.splitext(pdf_path)
        output_path = f"{filename}_bookmarked{ext}"
    
    logger.info(f"处理PDF文件: {pdf_path}")
    logger.info(f"输出到: {output_path}")
    
    reader = PdfReader(pdf_path)
    writer = PdfWriter()
    
    # 复制所有页面
    for page in reader.pages:
        writer.add_page(page)
    
    # 如果PDF已有书签且不强制重新检测，则复制现有书签
    existing_bookmarks = []
    outline_copied = False
    
    if reader.outline and not force_detect:
        logger.info("PDF已有书签，将保留现有书签")
        outline_copied = True
        
        def copy_bookmarks(outline, parent=None):
            for item in outline:
                if isinstance(item, list):
                    copy_bookmarks(item, parent)
                else:
                    if '/Title' in item:
                        title = item['/Title']
                        try:
                            # 尝试获取目标页码
                            page_num = reader.get_destination_page_number(item)
                            if parent is None:
                                parent_ref = writer.add_outline_item(title, page_num)
                            else:
                                parent_ref = writer.add_outline_item(title, page_num, parent)
                            
                            # 记录已添加的书签
                            existing_bookmarks.append((title, page_num))
                            
                            # 处理子书签
                            if '/First' in item:
                                child = item['/First']
                                while child:
                                    if '/Title' in child:
                                        child_title = child['/Title']
                                        try:
                                            child_page = reader.get_destination_page_number(child)
                                            writer.add_outline_item(child_title, child_page, parent_ref)
                                            existing_bookmarks.append((child_title, child_page))
                                        except:
                                            pass
                                    if '/Next' in child:
                                        child = child['/Next']
                                    else:
                                        break
                        except:
                            logger.warning(f"无法处理书签: {title}")
        
        if isinstance(reader.outline, list):
            copy_bookmarks(reader.outline)
        else:
            copy_bookmarks([reader.outline])
    
    # 如果需要自动检测标题并添加书签
    if auto_detect and (not outline_copied or force_detect):
        if force_detect and reader.outline:
            logger.info("强制模式：即使PDF有书签，也进行标题检测")
        
        logger.info("开始自动检测标题并添加书签")
        headings = extract_heading_candidates(pdf_path)
        
        if headings:
            logger.info(f"检测到 {len(headings)} 个可能的标题")
            
            # 根据标题层级和格式组织书签结构
            # 这里简单处理，实际应用中可能需要更复杂的算法
            current_level = {0: None}  # 每个层级的当前书签引用
            
            for title, page_num in headings:
                # 通过简单的规则估计标题层级
                level = 0
                
                # 判断层级的简单规则，可根据实际PDF格式调整
                if re.match(r'^\d+\.\d+\.\d+', title):  # 如 "1.1.1"
                    level = 2
                elif re.match(r'^\d+\.\d+', title):     # 如 "1.1"
                    level = 1
                elif re.match(r'^\d+\.', title):        # 如 "1."
                    level = 0
                elif re.match(r'^第\S+章', title):      # 如 "第一章"
                    level = 0
                elif re.match(r'^第\S+节', title):      # 如 "第一节"
                    level = 1
                
                # 添加书签
                if level == 0 or current_level.get(level-1) is None:
                    # 顶级书签或上级不存在时
                    ref = writer.add_outline_item(title, page_num-1)  # PyPDF2页码从0开始
                    current_level[level] = ref
                else:
                    # 添加到父书签下
                    parent = current_level[level-1]
                    ref = writer.add_outline_item(title, page_num-1, parent)
                    current_level[level] = ref
                
                # 清除更深层级的引用
                for k in list(current_level.keys()):
                    if k > level:
                        del current_level[k]
    
    # 保存修改后的PDF
    with open(output_path, "wb") as fp:
        writer.write(fp)
    
    logger.info(f"已成功添加书签并保存到: {output_path}")
    return output_path

def process_directory(directory_path, recursive=False, auto_detect=True, force_detect=False):
    """处理目录中的所有PDF文件"""
    logger.info(f"开始处理目录: {directory_path}")
    
    if not os.path.exists(directory_path):
        logger.error(f"目录不存在: {directory_path}")
        return
    
    all_pdfs = []
    
    # 收集所有PDF文件
    if recursive:
        for root, _, files in os.walk(directory_path):
            for file in files:
                if file.lower().endswith('.pdf'):
                    all_pdfs.append(os.path.join(root, file))
    else:
        all_pdfs = [os.path.join(directory_path, f) for f in os.listdir(directory_path) 
                   if os.path.isfile(os.path.join(directory_path, f)) and f.lower().endswith('.pdf')]
    
    all_pdfs.sort()  # 按文件名排序
    
    logger.info(f"找到 {len(all_pdfs)} 个PDF文件")
    
    # 处理每个PDF文件
    for pdf_path in all_pdfs:
        try:
            output_path = pdf_path.replace('.pdf', '_bookmarked.pdf')
            add_bookmarks_to_pdf(pdf_path, output_path, auto_detect, force_detect)
        except Exception as e:
            logger.error(f"处理文件 {pdf_path} 时出错: {str(e)}")
            logger.exception(e)

def main():
    parser = argparse.ArgumentParser(description='向PDF文件添加书签')
    parser.add_argument('path', help='PDF文件路径或包含PDF文件的目录路径')
    parser.add_argument('-r', '--recursive', action='store_true', help='如果指定的是目录，是否递归处理子目录')
    parser.add_argument('-o', '--output', help='输出文件路径(仅在处理单个文件时有效)')
    parser.add_argument('--no-auto-detect', action='store_true', help='不自动检测标题(仅保留已有书签)')
    parser.add_argument('-f', '--force', action='store_true', help='强制检测标题，即使PDF已有书签')
    
    args = parser.parse_args()
    
    if os.path.isdir(args.path):
        # 处理目录
        process_directory(args.path, args.recursive, not args.no_auto_detect, args.force)
    else:
        # 处理单个文件
        add_bookmarks_to_pdf(args.path, args.output, not args.no_auto_detect, args.force)

if __name__ == "__main__":
    main() 